#include <stdio.h>
#include <errno.h>
#include <unistd.h>
#include <stdlib.h>
#define _GNU_SOURCE
#include <string.h>
#include <cpuset.h>
#define _GNU_SOURCE
#include <getopt.h>

/* version string at help.c*/
extern char *csver;
extern void check_errno(const int line, const char *name, pid_t pid);
extern void show_all();
extern void show_help();

static pid_t pid = -1;
static int destcpu = -1;
static const char *procname = NULL;

static const char *optname  = NULL;
static const char *newname  = NULL;
static const char *cpulist  = NULL;
static const char *lbarea   = NULL;
static const char *execname    = NULL;

#define	forceflag 	 1<<1
#define migrateflag	 1<<2
#define killflag 	 1<<3
#define createflag	 1<<4
#define listflag	 1<<5
#define lbareaflag	 1<<6
#define delflag		 1<<7
#define bindflag 	 1<<8
#define newnameflag	 1<<9
#define jumpnameflag 1<<10
#define setflag 	 1<<11
#define cpulistflag	 1<<12
#define optnameflag	 1<<13
#define pidflag 	 1<<14
#define resetflag    1<<15
#define execflag      1<<16

static unsigned int optflag = 0;

static char *shortopt = ":abcde:fhklm:n:p:r:su:vx";
#define NOARG 	no_argument
#define HASARG 	required_argument
#define OTHER 	optional_argument
static struct option longopt[] = {
    {"unbind", NOARG, NULL, 0},
    {"unbindall", NOARG, NULL, 0},
    {"reset", NOARG, NULL, 2},
    {"migrate", HASARG, NULL, 'm'},
//              {"jump",        HASARG, NULL, 'j'},
    {"name", HASARG, NULL, 'n'},
    {"pid", HASARG, NULL, 'p'},
    {"cpu", HASARG, NULL, 'u'},
    {"rename", HASARG, NULL, 'r'},
    {"exec", HASARG, NULL, 'e'},
    {"setcpu", NOARG, NULL, 's'},
    {"exclude", NOARG, NULL, 'x'},
    {"bind", NOARG, NULL, 'b'},
    {"kill", NOARG, NULL, 'k'},
    {"list", NOARG, NULL, 'l'},
    {"help", NOARG, NULL, 'h'},
    {"del", NOARG, NULL, 'd'},
    {"create", NOARG, NULL, 'c'},
    {"force", NOARG, NULL, 'f'},
    {"version", NOARG, NULL, 'v'},
    {0, 0, 0, 0},
};

static void cs_create(void)
{
    struct cpuset *cs = (struct cpuset *) malloc(sizeof(*cs));
    memset(cs, 0, sizeof(*cs));
    strcpy(cs->name, optname);
    strcpy(cs->cpulist, cpulist);
    if (optflag & lbareaflag)
        cs->lbarea = 1;
    cpuset_create(cs, NULL);
    if (errno == ENOCPU)
		ck_error(1, errno, "%s: %s: ", procname, cpulist);
    ck_error(1, errno, "%s: %s: ", procname, optname);
}

static void cs_lbarea(void)
{
    switch (optflag) {
    case lbareaflag:		/* list lbarea */
			cpuset_showlbarea();
			exit(0);		/* exit normary */
			break;
    case lbareaflag | cpulistflag | optnameflag:	/* create a lbarea */
                        cs_create();
	break;
	/* del lbarea */
    case lbareaflag | delflag | optnameflag:
			cpuset_dellbarea(optname);
			ck_error(1, errno, "%s: %s: ", procname, optname);
			break;
    case lbareaflag | resetflag:
			cpuset_dellbarea_all();
			ck_error(1, errno, "%s: %s: ", procname, optname);
			break;
    }
}

static void cs_list(void)
{
    char buf[PATH_MAX];
    char cpulist[128];
    struct cpuset_pidlist *pl;
    struct cpuset_namelist *nl;
    int i;

    switch (optflag) {
    case listflag | optnameflag:
	// -l -n <name>
	if ((pl = cpuset_getpid_begin(optname)) != NULL) {
	    for (i = 0; i < pl->npids; ++i)
		fprintf(stdout, "%d\n", pl->pids[i]);
	    cpuset_getpid_end(pl);
	}
	ck_error(1, errno, "%s: %s: ", procname, optname);
	break;
    case listflag | pidflag:
	// -l -p <pid>  show which cpuset this pid resides in
	if (cpuset_getcpusetpath(pid, buf, PATH_MAX) != NULL)
	    fprintf(stdout, "%s", buf);
	if ((i = cpuset_latestcpu(pid)) > -1)
	    fprintf(stdout, "\t%d", i);
	if (cpuset_getcpulist(pid, cpulist) == 0)
	    fprintf(stdout, "\t%s\n", cpulist);
	ck_error(1, errno, "%s: %d: ", procname, pid);
	break;
    case listflag:
	// -l show all cpuset names and how many tasks 
	if ((nl = cpuset_getname_begin()) != NULL) {
	    for (i = 0; i < nl->num; ++i) {
		fprintf(stdout, "%s\t\t%s", nl->name[i], nl->cpulist[i]);
		if ((pl = cpuset_getpid_begin(nl->name[i])) != NULL)
		    fprintf(stdout, "\t\t%d\n", pl->npids);
		else
		    fprintf(stdout, "\t\t%d\n", 0);
		cpuset_getpid_end(pl);
	    }
	}
	ck_error(1, errno, "%s: ", procname);
    }
}

static void cs_rename(void)
{
    cpuset_rename(optname, newname);
    if (errno == EEXIST)
    	ck_error(1, errno, "%s: %s: ", procname, newname);
    else			//NOFILE
    	ck_error(1, errno, "%s: %s: ", procname, optname);
}

static void cs_setcpu(void)
{
    cpuset_setcpulist(optname, cpulist);
    if (errno == ENOCPU)
	ck_error(0, errno, "%s: %s: cpulist error!\n", procname, cpulist);
    else
	ck_error(1, errno, "%s: %s: ", procname, optname);
}

/* delete a cpuset , if -n <name> delete this cpuset ,if -p <pid> delete it 
   without force or not
   */
static void cs_del(void)
{
    int ret;
    if (optname == NULL) {
        if (pid < 0)		// no name ,no pid
        {
            fprintf(stdout, "cpuset: need to give a NAME or PID!\n");
            exit(EINVAL);
        }
        /* delete cpuset whitch PID resides in */
        cpuset_del_bypid(pid);
        if (errno == EINVAL)
            ck_error(0, errno, "%s: can't delete \"/\"\n", procname);
        ck_error(1, errno, "%s: %d: ", procname, pid);
    }
    /* delete cpuset by name */
    ret = cpuset_del_byname(optname, forceflag & optflag);
    if (ret < 0)
        ret = cpuset_dellbarea(optname);
    if (errno == EBUSY)
    	errno = ENOTEMPTY;
    if (errno == EINVAL)
    	ck_error(0, errno, "%s: can't delete \"/\"\n", procname);
    ck_error(1, errno, "%s: %s: ", procname, optname);
}
static int same_cpulist(pid_t pid, const char *cpulist)
{
	char pid_cpulist[128];
	if (cpuset_getcpulist(pid, pid_cpulist)<0) 
		return 0;
	if (strcmp(pid_cpulist, cpulist))
		return 0;
	return 1;
}

/* bind && ubind a process */
void cs_bind(void)
{
    char buf[PATH_MAX];
    struct cpuset cs;
    if (pid < 1)		// no pid gived
    {
	fprintf(stdout, "cpuset: no such pid:%d!\n", pid);
	exit(EINVAL);
    }
    /* bind PID to NAME */
    cpuset_getcpusetpath(pid, buf, PATH_MAX);
    if (optname != NULL) {
    	cpuset_getcpuset(optname, &cs);
	if (!strcmp(buf, optname) && same_cpulist(pid, cs.cpulist))
	    fprintf(stderr, "WARNNING: rebind %d to %s\n", pid, optname);
	else
	{
	    if (!strcmp(optname, cs.name))  //if same name, bind / first
                cpuset_bind(pid, "/");
	    cpuset_bind(pid, optname);
	}
	if (errno == ESRCH)
	    ck_error(1, errno, "%s: %d: ", procname, pid);
	ck_error(1, errno, "%s: %s: ", procname, optname);
    }
    /* give pid ,no cpuset name, unbind it */
    if (!strcmp(buf, "/"))

	fprintf(stderr, "WARNNING: rebind %d to %s\n", pid, "/");
    else
	cpuset_bind(pid, "/");
    ck_error(1, errno, "%s: %d: ", procname, pid);
}

static void cs_jump(void)
{
    if (!optname)
	ck_error(0, errno, "%s: need to give a name with --name  ",
		 procname);
    if (!strcmp(optname, "/"))
	ck_error(0, EINVAL, "%s: can't unbind \"/\"\n", procname);
    cpuset_bind_cpuset_tasks(optname, "/");
    ck_error(1, errno, "%s: %s : ", procname, optname);
}

static void cs_kill(void)
{
    if (optname == NULL) {
	fprintf(stderr, "cpuset: need to give a NAME!\n");
	exit(EINVAL);
    }
    if (!strcmp(optname, "/")) {
	fprintf(stderr, "cpuset: NO PERMITION TO KILL  \"/\"!\n");
	exit(EINVAL);
    }

    cpuset_kill(optname, forceflag & optflag);
    ck_error(1, errno, "%s: %s: ", procname, optname);
}

static void cs_migrate(void)
{
    cpuset_migrate(pid, destcpu);
    switch (errno) {
    case 0:			// no error
	break;
    case ESRCH:
	fprintf(stderr, "cpuset: %d: No such process\n", pid);
	break;
    case ENOCPU:
	fprintf(stderr, "cpuset: %d: Invalid cpu(may be not online)\n",
		destcpu);
	break;
    case ENOENT:
	fprintf(stderr, "cpuset: Kernel unsupport,USE K2OS!\n");
	break;
    case EBOUND:
	fprintf(stderr, "cpuset: Process %d bound by kernel\n", pid);
	break;
    default:
	fprintf(stderr, "cpuset: ERROR, errno: %d\n", errno);
    }
    exit(errno);		// exit on error
}

/* name is inherit from argv+optind */
static void cs_exec(char **name)
{	
	char *endptr;
	int cpu;

	if (optflag != (execflag | cpulistflag) && 
		 optflag != (execflag | optnameflag))
	{
			errno = EINVAL;
			goto err;
	}
	/* -e -u */
	if ( optflag == (execflag|cpulistflag) )
	{
		cpu = strtol(cpulist, &endptr, 10);
	    if (*endptr) 
		{
			fprintf(stderr, "%s: Cpu/Cpulist error.\n",procname);
			exit(EINVAL);
	    }
		cpuset_exec(name, cpu);
		goto err;
	}
	/* -e -n  bind self to cpuset first */
	if ( cpuset_bind(getpid(), optname) < 0)
		goto err;

	if (fork() == 0 )	/* child execute the program */
	{
		execvp(name[0],name);
		/* shouldn't come here */
		fprintf(stderr, "%s: %s\n", name[0], strerror(errno));
		exit(-1);
	}
err:
	ck_error(1, errno, "%s: ", procname);
}

int main(int argc, char **argv)
{
    int oc;
//      opterr = 0;
    int optindex;
    procname = argv[0];
    char *endptr;
    while ((oc =
	    getopt_long(argc, argv, shortopt, longopt, &optindex)) != -1) {
	//      printf("oc:%c, optarg:%s optind:%d argc:%d\n", oc, optarg, optind, argc);
	switch (oc) {
	case 0:
	    if (optindex == 0)	/* unbind */
		optflag |= bindflag;
	    else if (optindex == 1) {	/* unbindall */
		//      jumpname = optarg;  --unbindall -n <name>
		optflag |= jumpnameflag;
	    }
	    break;
	case 2:		// reset
	    optflag |= resetflag;
	    break;
	case 'x':
	    if (optarg)
		lbarea = optarg;
	    optflag |= lbareaflag;
	    break;
	case 'e':
	    optflag |= execflag;
		execname = optarg;
		cs_exec(argv + optind-1);
	    break;
	case 'k':
	    optflag |= killflag;
	    break;
	case 'b':
	    optflag |= bindflag;
	    break;
	case 'm':
	    optflag |= migrateflag;
	    destcpu = strtol(optarg, &endptr, 10);
	    if (*endptr) {
		fprintf(stderr, "%s: Please give the right CPU.\n",
			procname);
		exit(EINVAL);
	    }
	    break;
	case 'n':
	    optname = optarg;
	    optflag |= optnameflag;
	    break;
	case 'p':
	    optflag |= pidflag;
	    pid = (pid_t) strtol(optarg, &endptr, 10);
	    if (*endptr) {
		fprintf(stderr, "%s: Please give the right PID.\n",
			procname);
		exit(EINVAL);
	    }
	    break;
	case 'u':
	    cpulist = optarg;
	    optflag |= cpulistflag;
	    break;
	case 'r':
	    newname = optarg;
	    optflag |= newnameflag;
	    break;
	case 's':
	    optflag |= setflag;
	    break;
	case 'l':
	    optflag |= listflag;
	    break;
	case 'h':
	    show_help();
	    break;
	case 'd':
	    optflag |= delflag;
	    break;
	case 'c':
	    optflag |= createflag;
	    break;
	case 'f':
	    optflag |= forceflag;
	    break;
/*				case 'j':
						jumpname = optarg;;
						optflag |= jumpnameflag;
								break;
*/ case 'v':
	    fprintf(stdout, "cpuset %s\n", csver);
	    exit(0);
	    break;
	case ':':
	    fprintf(stderr, "cpuset: %s: Need more argument\n",
		    argv[optind - 1]);
	    errno = -EINVAL;
	    goto exit;
	case '?':
	default:
	    fprintf(stderr, "cpuset: %s: Invalid option\n",
		    argv[optind - 1]);
	    errno = -EINVAL;
	    goto exit;

	}

    }
/* this should be fixed to support many -p <pid1,pid2,pid3...> */
    if (optind < argc) {
	fprintf(stderr, "cpuset: Not option value: ");
	while (optind < argc)
	    fprintf(stderr, "%s ", argv[optind++]);
	fprintf(stderr, "\n");
	exit(EINVAL);
    }

    switch( optflag )
    {
        case createflag | optnameflag | cpulistflag:
        case lbareaflag | createflag |cpulistflag | optnameflag :/* create a lbarea */
            cs_create();
            break;
        case newnameflag | optnameflag:
            cs_rename();
            break;
        case setflag | optnameflag | cpulistflag:
            cs_setcpu();
            break;
        case delflag | optnameflag:	/* -d -n <name> */
        case delflag | optnameflag | forceflag: /* -d -n <name> -f */
        case delflag | pidflag:/* -d -p <pid> */
            cs_del();
            break;
        case bindflag | pidflag:	/* unbind it */
        case bindflag | optnameflag | pidflag:
            cs_bind();
            break;
        /* -j <des_name> -n <source_name> */
        /* -j <unbind all name> */
        case jumpnameflag | optnameflag:
            cs_jump();
            break;
        case migrateflag | pidflag:	/* -m -p <pid> */
            cs_migrate();
            break;
        case killflag | optnameflag:	/* -k -n <name> */
        case killflag | optnameflag | forceflag:	/* -k -n <name> -f */
            cs_kill();
            break;
        case lbareaflag:	/* list lbarea */
        case lbareaflag | cpulistflag | optnameflag:	/* create a lbarea */
        case lbareaflag | delflag | optnameflag:	/* del lbarea */
	    case lbareaflag | resetflag:	/* del lbarea */
            cs_lbarea();
            break;
        case listflag:	/* -l */
        case listflag | pidflag:/* -l -p <pid> */
        case listflag | optnameflag:/* -l -n <name> */
            cs_list();
            break;
        default:
            errno = EINVAL;
    /* nothing to do, show help */
            show_help();
    }
  exit:
    exit(errno);
}
